# file: /data/pyworks/fw_paddleocr_web/app/core/ocr_processor.py

import asyncio
from typing import Dict, Any

from app.utils.logger import get_logger


class OCRProcessor:
    def __init__(self):
        self.logger = get_logger(__name__)
        # 初始化所有模块为 None
        self.text_recognition_module = None
        self.layout_analysis_module = None
        self.table_structure_recognition = None
        self.formula_recognition_module = None
        self.doc_orientation_module = None
        self.seal_text_detection_module = None
        self.chart_parser_module = None
        self.vlm_module = None

        # 标记各模块是否已初始化
        self._modules_initialized = {
            'text': False,
            'layout': False,
            'table': False,
            'formula': False,
            'orientation': False,
            'seal': False,
            'chart': False,
            'vlm': False
        }

    def _get_text_module(self):
        """获取文本识别模块（延迟加载）"""
        if not self._modules_initialized['text']:
            from app.modules.text_recognition import TextRecognitionModule
            self.text_recognition_module = TextRecognitionModule()
            self._modules_initialized['text'] = True
        return self.text_recognition_module

    def _get_layout_module(self):
        """获取版面分析模块（延迟加载）"""
        if not self._modules_initialized['layout']:
            from app.modules.layout_analysis import LayoutAnalysisModule
            self.layout_analysis_module = LayoutAnalysisModule()
            self._modules_initialized['layout'] = True
        return self.layout_analysis_module

    def _get_table_module(self):
        """获取表格结构识别模块（延迟加载）"""
        if not self._modules_initialized['table']:
            from app.modules.table_structure_recognition import TableStructureRecognitionModule
            self.table_structure_recognition = TableStructureRecognitionModule()
            self._modules_initialized['table'] = True
        return self.table_structure_recognition

    def _get_formula_module(self):
        """获取公式识别模块（延迟加载）"""
        if not self._modules_initialized['formula']:
            from app.modules.formula_recognition import FormulaRecognitionModule
            self.formula_recognition_module = FormulaRecognitionModule()
            self._modules_initialized['formula'] = True
        return self.formula_recognition_module

    def _get_orientation_module(self):
        """获取文档方向分类模块（延迟加载）"""
        if not self._modules_initialized['orientation']:
            from app.modules.doc_orientation import DocumentOrientationModule
            self.doc_orientation_module = DocumentOrientationModule()
            self._modules_initialized['orientation'] = True
        return self.doc_orientation_module

    def _get_seal_text_module(self):
        """获取印章文本检测模块（延迟加载）"""
        if not self._modules_initialized['seal']:
            from app.modules.seal_text_detection import SealTextDetectionModule
            self.seal_text_detection_module = SealTextDetectionModule()
            self._modules_initialized['seal'] = True
        return self.seal_text_detection_module

    def _get_chart_module(self):
        """获取图表解析模块（延迟加载）"""
        if not self._modules_initialized['chart']:
            from app.modules.chart_parser import ChartParserModule
            self.chart_parser_module = ChartParserModule()
            self._modules_initialized['chart'] = True
        return self.chart_parser_module

    def _get_vlm_module(self):
        """获取视觉语言模型模块（延迟加载）"""
        if not self._modules_initialized['vlm']:
            from app.modules.vlm import VisionLanguageModule
            self.vlm_module = VisionLanguageModule()
            self._modules_initialized['vlm'] = True
        return self.vlm_module

    async def process_text_recognition(
            self,
            image_data: bytes,
            language: str = "ch",
            detail: bool = False
    ) -> Dict[str, Any]:
        """处理文本识别请求"""
        self.logger.info(f"Processing text recognition, language: {language}")
        module = self._get_text_module()  # 获取模块实例
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            module.recognize,
            image_data,
            language,
            detail
        )
        return result

    async def process_layout_analysis(
            self,
            image_data: bytes,
            detect_tables: bool = True,
            detect_formulas: bool = True
    ) -> Dict[str, Any]:
        """处理版面分析请求"""
        self.logger.info("Processing layout analysis")
        module = self._get_layout_module()  # 获取模块实例
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            module.analyze,
            image_data,
            detect_tables,
            detect_formulas
        )
        return result

    async def process_table_structure_recognition(
            self,
            image_data: bytes,
            structure_only: bool = False
    ) -> Dict[str, Any]:
        """处理表格结构识别请求"""
        self.logger.info("Processing table recognition")
        module = self._get_table_module()  # 获取模块实例
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            module.recognize,
            image_data,
            structure_only
        )
        return result

    async def process_formula_recognition(
            self,
            image_data: bytes,
            format_type: str = "latex"
    ) -> Dict[str, Any]:
        """处理公式识别请求"""
        self.logger.info(f"Processing formula recognition, format: {format_type}")
        module = self._get_formula_module()  # 获取模块实例
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            module.recognize,
            image_data,
            format_type
        )
        return result

    async def process_doc_orientation(
            self,
            image_data: bytes
    ) -> Dict[str, Any]:
        """处理文档方向分类请求"""
        self.logger.info("Processing document orientation classification")
        module = self._get_orientation_module()  # 获取模块实例
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            module.classify,
            image_data
        )
        return result

    async def process_seal_text_detection(
            self,
            image_data: bytes
    ) -> Dict[str, Any]:
        """处理印章文本检测请求"""
        self.logger.info("Processing seal text detection")
        module = self._get_seal_text_module()  # 获取模块实例
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            module.detect,
            image_data
        )
        return result

    async def process_chart_parsing(
            self,
            image_data: bytes,
            prompt: str = ""
    ) -> Dict[str, Any]:
        """处理图表解析请求"""
        self.logger.info("Processing chart parsing")
        module = self._get_chart_module()  # 获取模块实例
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            module.parse,
            image_data
        )
        return result

    async def process_vision_language_analysis(
            self,
            image_data: bytes,
            prompt: str = "Describe this image"
    ) -> Dict[str, Any]:
        """处理视觉语言模型分析请求"""
        self.logger.info("Processing vision language analysis")
        module = self._get_vlm_module()  # 获取模块实例
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            module.analyze,
            image_data,
            prompt
        )
        return result
